from typing import Dict, Callable, Any
from functools import wraps
from concurrent.futures import ThreadPoolExecutor
from models.enums import EventMessageType
from dataprocess.handler.RagEventHandler import RagEventHandler
from dataprocess.handler.FileEventHandler import FileEventHandler
import asyncio
import threading
from Logger import Logger

logger = Logger.get_logger(__name__)

class RagTaskManger:
    _executor = ThreadPoolExecutor(max_workers=20)
    _handlers: Dict[str, Callable] = {}
    _loop = None
    _loop_thread = None
    _loop_ready = threading.Event()

    @classmethod
    def _init_event_loop(cls):
        """Initialize event loop in a separate thread"""
        if cls._loop is None:
            def run_event_loop():
                try:
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                    cls._loop = loop
                    cls._loop_ready.set()  # Signal that loop is ready
                    loop.run_forever()
                except Exception as e:
                    logger.error(f"Event loop initialization error: {str(e)}")
                    cls._loop_ready.set()  # Set event even on error

            cls._loop_thread = threading.Thread(target=run_event_loop, daemon=True)
            cls._loop_thread.start()
            # Wait for loop to be ready
            cls._loop_ready.wait(timeout=5.0)

    @classmethod
    def _run_coroutine(cls, coro):
        """Run coroutine in the event loop"""
        try:
            if cls._loop is None:
                cls._init_event_loop()
            
            if cls._loop is None:
                raise RuntimeError("Failed to initialize event loop")

            future = asyncio.run_coroutine_threadsafe(coro, cls._loop)
            # Optional: You can add a callback to handle completion
            future.add_done_callback(
                lambda f: logger.info("Task completed" if not f.exception() else f"Task failed: {f.exception()}")
            )
        except Exception as e:
            logger.error(f"Error running coroutine: {str(e)}")
            raise

    @classmethod
    def start_task(cls, user_message: dict):
        handler_name = user_message['handler']
        try:
            match handler_name:
                case EventMessageType.RAG.value:
                    cls._run_coroutine(RagEventHandler.handle(user_message))
                case EventMessageType.FILE.value:
                    cls._run_coroutine(FileEventHandler.handle(user_message))
                case _:
                    raise Exception("任务启动失败，未找到合适的任务处理器！")
        except Exception as e:
            logger.error(f"Error starting task: {str(e)}")
            raise  # Re-raise to let caller handle it


